import csv
import os
from .enums import PastTime, DeviceName
from db import DataBase

from typing import List


def gen_csv(client: DataBase, file_name: str, time: PastTime = PastTime.ONE_HOUR):
    """生成csv数据集"""
    file_path = os.path.join(os.path.abspath("."), "data", file_name)
    with open(file_path, 'w', newline='') as f:
        writer = csv.writer(f)
        points = client.query(DeviceName.SUMMARY, time)
        try:
            header = list(points[0].keys())
        except IndexError:
            raise IndexError("DB have no data")
        header = header[1:]  # 越过time这个key
        writer.writerow(header)
        for point in points:
            data = list(point.values())[1:]
            writer.writerow(data)


def read_csv(file_name: str) -> List[List[float]]:
    """读取csv文件"""
    file_path = os.path.join(os.path.abspath("."), "data", file_name)
    if not os.path.exists(file_path):
        raise FileNotFoundError("CSV file: '{}' not found!".format(file_path))

    rsp = []
    with open(file_path, 'r') as f:
        reader = csv.reader(f)
        for index, row in enumerate(reader):
            if index == 0:
                continue
            row = list(map(float, row))
            rsp.append(row)

    return rsp
